import os

from transformers import BertTokenizerFast
import pickle
from tqdm import tqdm

current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


class DataPreprocess:
    """
    数据预处理
    """

    def __init__(self):
        self.tokenizer = None

    def preprocess(self, txt_path, pkl_path):
        """
        对原始语料进行tokenize，将每段对话处理成如下形式："[CLS]utterance1[SEP]utterance2[SEP]utterance3[SEP]"
        :param txt_path:
        :param pkl_path:
        :return:
        """
        """
        BertTokenizerFast 是 Hugging Face transformers 库提供的 高性能 BERT 分词器，
        基于 Rust 实现，比标准 BertTokenizer 速度更快，适合处理大批量文本。
        1. 核心特性
        特性	说明
        基于 Rust 实现 比 Python 实现的 BertTokenizer 快 10-100 倍
        支持并行化 批量分词时自动多线程处理
        兼容性 与 BertTokenizer API 完全一致，可直接替换
        子词分词 使用 BPE（Byte-Pair Encoding）算法处理未登录词
        """
        #  初始化tokenizer，使用BertTokenizerFast从预训练的中文Bert模型（bert-base-chinese）创建一个tokenizer对象

        #  vocab_size 词表的大小（即分词器能识别的唯一 token 数量）
        # print(f'唯一token数量 tokenizer.vocab_size-->{tokenizer.vocab_size}')
        sep_id = self.tokenizer.sep_token_id  # 获取分隔符[SEP]的token ID
        cls_id = self.tokenizer.cls_token_id  # 获取起始符[CLS]的token ID
        # print(f'sep_id-->{sep_id}')
        print(f'cls_id-->{cls_id}')
        # 读取数据集
        with open(os.path.join(current_dir, txt_path), 'r', encoding="utf-8") as f:
            data = f.read()  # 以UTF-8编码读取文件内容
        # Windows 换行符：\r\n
        # Linux/macOS 换行符：\n
        if "\r\n" in data:
            data_split = data.split("\r\n\r\n")
        else:
            data_split = data.split("\n\n")

        # print(len(data_split))  # 打印对话段落数量

        # 开始进行tokenize
        # 保存所有的对话数据,每条数据的格式为："[CLS]seq1[SEP]seq2[SEP]seq3[SEP]"

        # 记录所有对话tokenize分词之后的长度，用于统计中位数与均值
        dialogue_len = []
        dialogue_list = []  # 记录所有对话
        # # #
        for index, text_info in enumerate(tqdm(data_split)):
            # 按行切割
            if "\r\n" in text_info:
                sentences = text_info.split("\r\n")
            else:
                sentences = text_info.split("\n")

            input_ids = [cls_id]  # 每个dialogue以[CLS]seq1[sep]seq2[sep]开头
            for sentence in sentences:
                # print(f'sentence-->{sentence}')
                # 将每个句子进行tokenize，并将结果拼接到input_ids列表中
                sentence_id = self.tokenizer.encode(sentence, add_special_tokens=False)
                # print(f'sentence_id-->{sentence_id}')
                input_ids += sentence_id
                # 每个seq之后添加[SEP]，表示seqs会话结束
                input_ids.append(sep_id)
            # 将对话的tokenize后的长度添加到对话长度列表中
            # print(input_ids)
            dialogue_len.append(len(input_ids))
            # 将tokenize后的对话添加到对话列表中
            dialogue_list.append(input_ids)

        with open(os.path.join(current_dir, pkl_path), "wb") as pkl_f:
            pickle.dump(dialogue_list, pkl_f)


if __name__ == '__main__':
    data_preprocess = DataPreprocess()
    train_txt_path = 'data/medical_train.txt'
    train_pkl_path = 'data/medical_train.pkl'
    data_preprocess.preprocess(train_txt_path, train_pkl_path)
